'''
Created on Jul 26, 2011

@author: peterd
'''

import os

import networkx as nx

from mhd.models import MHD_Stop, MHD_Link, MHD_Link_Direction, MHD_Link_Stop 
from mhd.locator import MHD_StopLocator, InvalidLatLng
from mhd.exceptions import NoSuchStop

class MHD_GraphGenerator:
    def __init__(self):
        self.mhdG_wrkdays_school_ = nx.MultiDiGraph(ver='1.0', type='wrkdays_school')
        self.mhdG_wrkdays_noschool_ = nx.MultiDiGraph(ver='1.0', type='wrkdays_noschool')
        self.mhdG_weekend_ = nx.MultiDiGraph(ver='1.0', type='weekend')
        self.stopLocator_ = MHD_StopLocator()                

    
    def generate(self):        
        # first create stop nodes and walking edges
        for stop in MHD_Stop.objects.all():
            print "adding stop %s to graphs" % stop.id
            self.mhdG_wrkdays_school_.add_node(stop.id, id = stop.id, name = stop.name, lat = stop.lat, lng = stop.lng, zoid = stop.zoid, type = stop.stop_type_id, links = [])
            self.mhdG_wrkdays_noschool_.add_node(stop.id, id = stop.id, name = stop.name, lat = stop.lat, lng = stop.lng, zoid = stop.zoid, type = stop.stop_type_id, links = [])
            self.mhdG_weekend_.add_node(stop.id, id = stop.id, name = stop.name, lat = stop.lat, lng = stop.lng, zoid = stop.zoid, type = stop.stop_type_id, links = [])
            
            # consider only neighbouring stops within 250m to constitute candidate walking edges
            try:
                for nearbyStop in self.stopLocator_.find_nearby_stops_by_latlng(stop.lat, stop.lng, 250):   
                    if not set(stop.getLinks()) & set(nearbyStop[0].getLinks()):                                            
                        self.mhdG_wrkdays_school_.add_edge(stop.id, nearbyStop[0].id, type = 'w', weight = nearbyStop[1] / 60)
                        self.mhdG_wrkdays_noschool_.add_edge(stop.id, nearbyStop[0].id, type = 'w', weight = nearbyStop[1] / 60)
                        self.mhdG_weekend_.add_edge(stop.id, nearbyStop[0].id, type = 'w', weight = nearbyStop[1] / 60)
            except InvalidLatLng:
                continue
                    
        for link in MHD_Link.objects.all():  
            #if not link.name in ['20', '22', '23', '34']:
            #if not link.name in ['34']:                
            #    continue           
            graphMap = {}
            graphMap[self.mhdG_wrkdays_school_] = link.operatesWorkDays()
            graphMap[self.mhdG_wrkdays_noschool_] = link.operatesSchoolHolidays()
            graphMap[self.mhdG_weekend_] = link.operatesWeekends()
            
            for g, linkOperates in graphMap.iteritems():
                if linkOperates:
                    print "adding link %s to graph %s" % (link.name, g.graph['type'])             
                    for direction in MHD_Link_Direction.objects.filter(link = link):                
                        prevStop = None
                        linkStopBase = None
                        stopNum = 1
                        isFirst = True
                        for linkStop in MHD_Link_Stop.objects.filter(link = link, direction = direction).order_by('stop_number'):                            
                            if prevStop:
                                try:
                                    if prevStop.id != linkStop.stop.id:
                                        g.add_edge(prevStop.id, linkStop.stop.id, type = 'd', dir = direction.id, weight = link.getDelta(linkStopBase, prevStop, linkStop.stop), key = link.name)
                                except NoSuchStop:
                                    pass
                                isFirst = False
                            else:
                                linkStopBase = linkStop

                            g.node[linkStop.stop.id]['links'].append({'name': linkStop.link.name, 'type': linkStop.link.link_type.id, 'dir': direction.id, 'stop_num': stopNum, 'first': isFirst, 'last': False})                                                           
                                                        
                            prevStop = linkStop.stop
                            stopNum = stopNum + 1
                                                    
                        for l in g.node[prevStop.id]['links']:
                            if l['name'] == link.name and l['dir'] == direction.id:
                                l['last'] = True 
                                break                                                                    
        
        if os.getcwd().find('/mhd') != -1:
            nx.write_gpickle(self.mhdG_wrkdays_school_, '../graphs/mhdg_workdays_school.pickle')
            nx.write_gpickle(self.mhdG_wrkdays_noschool_, '../graphs/mhdg_workdays_noschool.pickle')
            nx.write_gpickle(self.mhdG_weekend_, '../graphs/mhdg_weekend.pickle')
        else:
            nx.write_gpickle(self.mhdG_wrkdays_school_, 'graphs/mhdg_workdays_school.pickle')
            nx.write_gpickle(self.mhdG_wrkdays_noschool_, 'graphs/mhdg_workdays_noschool.pickle')
            nx.write_gpickle(self.mhdG_weekend_, 'graphs/mhdg_weekend.pickle')
                
def main():
    mhdGraphGen = MHD_GraphGenerator()
    mhdGraphGen.generate()
    
        
if __name__ == '__main__':
    main()
